// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <cstdint>
#include <pollux/testing/functions/prestosql/function_base_test.h>

using namespace kumo::pollux;
using namespace kumo::pollux::exec;
using namespace kumo::pollux::functions::test;

namespace {

class MapKeysAndValuesTest : public FunctionBaseTest {
 protected:
  void checkResult(
      MapVectorPtr map_vector,
      VectorPtr mapKeys,
      vector_size_t row,
      DecodedVector& decodedArray,
      const ArrayVector* result,
      VectorPtr resultElements) {
    auto size = map_vector->sizeAt(row);
    EXPECT_EQ(size, result->sizeAt(decodedArray.index(row))) << "at " << row;
    for (vector_size_t j = 0; j < size; j++) {
      EXPECT_TRUE(mapKeys->equal_value_at(
          resultElements->wrapped_vector(),
          map_vector->offsetAt(row) + j,
          result->offsetAt(decodedArray.index(row)) + j));
    }
  }

  void testMap(
      const std::string& expression,
      std::function<vector_size_t(vector_size_t /* row */)> sizeAt,
      std::function<bool(vector_size_t /*row */)> is_null_at,
      std::function<VectorPtr(MapVectorPtr /*map_vector*/)> expectedFunc) {
    auto map_vector = make_map_vector<int32_t, int64_t>(
        numRows_, sizeAt, keyAt, value_at, is_null_at);

    auto result = evaluate(expression, make_row_vector({map_vector}));
    DecodedVector decodedArray(*result);
    auto baseArray = decodedArray.base()->as<ArrayVector>();
    ASSERT_EQ(result->type_kind(), TypeKind::ARRAY);

    auto expected = expectedFunc(map_vector);

    auto baseArrayElements = baseArray->elements();
    ASSERT_TRUE(baseArrayElements->type()->kindEquals(expected->type()));

    EXPECT_EQ(numRows_, decodedArray.size());
    for (vector_size_t i = 0; i < numRows_; ++i) {
      EXPECT_EQ(decodedArray.is_null_at(i), map_vector->is_null_at(i)) << "at " << i;
      if (!map_vector->is_null_at(i)) {
        checkResult(
            map_vector, expected, i, decodedArray, baseArray, baseArrayElements);
      }
    }
  }

  void testMapPartiallyPopulated(
      const std::string& expression,
      std::function<vector_size_t(vector_size_t /* row */)> sizeAt,
      std::function<bool(vector_size_t /*row */)> is_null_at,
      std::function<VectorPtr(MapVectorPtr /*map_vector*/)> expectedFunc) {
    auto a = make_map_vector<int32_t, int64_t>(
        numRows_, sizeAt, keyAt, value_at, is_null_at);
    auto b = make_map_vector<int32_t, int64_t>(
        numRows_, sizeAt, value_at, keyAt, is_null_at);
    auto c =
        make_flat_vector<int32_t>(numRows_, [](vector_size_t i) { return i; });

    auto result = evaluate(expression, make_row_vector({c, a, b}));
    DecodedVector decodedArray(*result);
    auto baseArray = decodedArray.base()->as<ArrayVector>();
    ASSERT_EQ(result->type_kind(), TypeKind::ARRAY);

    auto aExpected = expectedFunc(a);
    auto bExpected = expectedFunc(b);

    auto baseArrayElements = baseArray->elements();
    ASSERT_TRUE(baseArray->type()->childAt(0)->kindEquals(aExpected->type()));

    EXPECT_EQ(numRows_, decodedArray.size());
    for (vector_size_t i = 0; i < numRows_; ++i) {
      if (i % 2 == 0) {
        EXPECT_EQ(decodedArray.is_null_at(i), a->is_null_at(i)) << "at " << i;
        if (!decodedArray.is_null_at(i)) {
          checkResult(
              a, aExpected, i, decodedArray, baseArray, baseArrayElements);
        }
      } else {
        EXPECT_EQ(decodedArray.is_null_at(i), b->is_null_at(i)) << "at " << i;
        if (!decodedArray.is_null_at(i)) {
          checkResult(
              b, bExpected, i, decodedArray, baseArray, baseArrayElements);
        }
      }
    }
  }

  static inline int32_t keyAt(vector_size_t idx) {
    return idx;
  }

  static inline int64_t value_at(vector_size_t idx) {
    return (idx + 1) * 11;
  }

 private:
  const vector_size_t numRows_ = 100;
};

class MapKeysTest : public MapKeysAndValuesTest {
 protected:
  void testMapKeys(
      std::function<vector_size_t(vector_size_t /* row */)> sizeAt,
      std::function<bool(vector_size_t /*row */)> is_null_at) {
    testMap("map_keys(C0)", sizeAt, is_null_at, mapKeys);
  }

  void testConstantMapKeys(
      std::function<vector_size_t(vector_size_t /* row */)> sizeAt,
      std::function<bool(vector_size_t /*row */)> is_null_at) {
    testMap("map_keys(C0)", sizeAt, is_null_at, mapKeys);
  }

  void testMapKeysPartiallyPopulated(
      std::function<vector_size_t(vector_size_t /* row */)> sizeAt,
      std::function<bool(vector_size_t /*row */)> is_null_at) {
    testMapPartiallyPopulated(
        "if(C0 % 2 = 0, map_keys(C1), map_keys(C2))",
        sizeAt,
        is_null_at,
        mapKeys);
  }

 private:
  static VectorPtr mapKeys(MapVectorPtr map_vector) {
    return map_vector->mapKeys();
  }
};

class MapValuesTest : public MapKeysAndValuesTest {
 protected:
  void testMapValues(
      std::function<vector_size_t(vector_size_t /* row */)> sizeAt,
      std::function<bool(vector_size_t /*row */)> is_null_at) {
    testMap("map_values(C0)", sizeAt, is_null_at, mapValues);
  }

  void testConstantMapValues(
      std::function<vector_size_t(vector_size_t /* row */)> sizeAt,
      std::function<bool(vector_size_t /*row */)> is_null_at) {
    testMap("map_values(C0)", sizeAt, is_null_at, mapValues);
  }

  void testMapValuesPartiallyPopulated(
      std::function<vector_size_t(vector_size_t /* row */)> sizeAt,
      std::function<bool(vector_size_t /*row */)> is_null_at) {
    testMapPartiallyPopulated(
        "if(C0 % 2 = 0, map_values(C1), map_values(C2))",
        sizeAt,
        is_null_at,
        mapValues);
  }

 private:
  static VectorPtr mapValues(MapVectorPtr map_vector) {
    return map_vector->mapValues();
  }
};

} // namespace

TEST_F(MapKeysTest, noNulls) {
  auto sizeAt = [](vector_size_t row) { return row % 7; };
  testMapKeys(sizeAt, nullptr);
}

TEST_F(MapKeysTest, someNulls) {
  auto sizeAt = [](vector_size_t row) { return row % 7; };
  testMapKeys(sizeAt, null_every(5));
}

TEST_F(MapKeysTest, allNulls) {
  auto sizeAt = [](vector_size_t row) { return row % 7; };
  testConstantMapKeys(sizeAt, null_every(1));
}

TEST_F(MapKeysTest, partiallyPopulatedNoNulls) {
  auto sizeAt = [](vector_size_t /* row */) { return 1; };
  testMapKeysPartiallyPopulated(sizeAt, nullptr);
}

TEST_F(MapKeysTest, partiallyPopulatedSomeNulls) {
  auto sizeAt = [](vector_size_t /* row */) { return 1; };
  testMapKeysPartiallyPopulated(sizeAt, null_every(5));
}

TEST_F(MapKeysTest, constant) {
  vector_size_t size = 1'000;
  auto data = make_map_vector<int64_t, int64_t>({
      {
          {0, 0},
          {1, 10},
          {2, 20},
          {3, 30},
      },
      {
          {4, 40},
          {5, 50},
      },
      {
          {6, 60},
      },
  });

  auto evaluateConstant = [&](vector_size_t row, const VectorPtr& vector) {
    return evaluate(
        "map_keys(c0)",
        make_row_vector({BaseVector::wrap_in_constant(size, row, vector)}));
  };

  auto result = evaluateConstant(0, data);
  auto expected = make_constant_array<int64_t>(size, {0, 1, 2, 3});
  test::assertEqualVectors(expected, result);

  result = evaluateConstant(1, data);
  expected = make_constant_array<int64_t>(size, {4, 5});
  test::assertEqualVectors(expected, result);

  result = evaluateConstant(2, data);
  expected = make_constant_array<int64_t>(size, {6});
  test::assertEqualVectors(expected, result);
}

TEST_F(MapKeysTest, unknown) {
  auto keys = make_flat_vector<UnknownValue>({}, UNKNOWN());
  auto values = make_flat_vector<UnknownValue>({}, UNKNOWN());
  auto map = make_map_vector({0, 0, 0}, keys, values);

  auto result = evaluate("map_keys(c0)", make_row_vector({map}));
  auto expected = make_array_vector({0, 0, 0}, keys);

  test::assertEqualVectors(expected, result);
}

TEST_F(MapValuesTest, noNulls) {
  auto sizeAt = [](vector_size_t row) { return row % 7; };
  testMapValues(sizeAt, nullptr);
}

TEST_F(MapValuesTest, someNulls) {
  auto sizeAt = [](vector_size_t row) { return row % 7; };
  testMapValues(sizeAt, null_every(5));
}

TEST_F(MapValuesTest, allNulls) {
  auto sizeAt = [](vector_size_t row) { return row % 7; };
  testConstantMapValues(sizeAt, null_every(1));
}

TEST_F(MapValuesTest, partiallyPopulatedNoNulls) {
  auto sizeAt = [](vector_size_t /* row */) { return 1; };
  testMapValuesPartiallyPopulated(sizeAt, nullptr);
}

TEST_F(MapValuesTest, partiallyPopulatedSomeNulls) {
  auto sizeAt = [](vector_size_t /* row */) { return 1; };
  testMapValuesPartiallyPopulated(sizeAt, null_every(5));
}

TEST_F(MapValuesTest, constant) {
  vector_size_t size = 1'000;
  auto data = make_map_vector<int64_t, int64_t>({
      {
          {0, 0},
          {1, 10},
          {2, 20},
          {3, 30},
      },
      {
          {4, 40},
          {5, 50},
      },
      {
          {6, 60},
      },
  });

  auto evaluateConstant = [&](vector_size_t row, const VectorPtr& vector) {
    return evaluate(
        "map_values(c0)",
        make_row_vector({BaseVector::wrap_in_constant(size, row, vector)}));
  };

  auto result = evaluateConstant(0, data);
  auto expected = make_constant_array<int64_t>(size, {0, 10, 20, 30});
  test::assertEqualVectors(expected, result);

  result = evaluateConstant(1, data);
  expected = make_constant_array<int64_t>(size, {40, 50});
  test::assertEqualVectors(expected, result);

  result = evaluateConstant(2, data);
  expected = make_constant_array<int64_t>(size, {60});
  test::assertEqualVectors(expected, result);
}

TEST_F(MapValuesTest, noDictionaryReceived) {
  // This tests make sure that expression evaluation does the peeling before
  // calling the function map_keys, even if it is performed before on the inputs
  // because they are all constants.
  vector_size_t size = 10;
  BufferPtr indices = AlignedBuffer::allocate<vector_size_t>(size, pool_.get());
  auto rawIndices = indices->asMutable<vector_size_t>();

  for (size_t i = 0; i < size; ++i) {
    rawIndices[i] = size - 1 - i;
  }

  auto map_vector = make_map_vector<int64_t, int64_t>(
      {{{1, 1}, {2, 2}}, {{100, 11}, {200, 1}}});

  auto constantMap = BaseVector::wrap_in_constant(2, 1, map_vector);

  auto result = evaluate(
      "map_keys(subscript(array_constructor(c0),1))",
      make_row_vector({constantMap}));
  auto expected = make_array_vector<int64_t>({{100, 200}, {100, 200}});
  test::assertEqualVectors(expected, result);
}

TEST_F(MapValuesTest, unknown) {
  auto keys = make_flat_vector<UnknownValue>({}, UNKNOWN());
  auto values = make_flat_vector<UnknownValue>({}, UNKNOWN());
  auto map = make_map_vector({0, 0, 0}, keys, values);

  auto result = evaluate("map_values(c0)", make_row_vector({map}));
  auto expected = make_array_vector({0, 0, 0}, values);

  test::assertEqualVectors(expected, result);
}
